import os
import xml.etree.ElementTree as ET
from PIL import Image


def prettify(element, indent='  '):
  """递归地遍历XML元素，添加缩进和换行以美化输出。"""
  queue = [(0, element)]  # (level, element)
  while queue:
    level, element = queue.pop(0)
    children = [(level + 1, child) for child in list(element)]
    if children:
      element.text = '\n' + indent * (level + 1)  # 子元素前的缩进
    if queue:
      element.tail = '\n' + indent * queue[0][0]  # 同级元素之间的缩进
    else:
      element.tail = '\n' + indent * (level - 1)  # 结尾的缩进
    queue[0:0] = children  # 将子元素插入队列

def create_xml_file(img_path, txt_path, xml_path, classes):
  # 从图像路径中提取文件名和尺寸
  img_file = Image.open(img_path)
  width, height = img_file.size
  img_file.close()

  # 创建XML文件基本结构
  annotation = ET.Element('annotation')
  ET.SubElement(annotation, 'folder').text = 'Images'
  ET.SubElement(annotation, 'filename').text = os.path.basename(img_path)
  ET.SubElement(annotation, 'path').text = img_path
  source = ET.SubElement(annotation, 'source')
  ET.SubElement(source, 'database').text = 'Unknown'
  size = ET.SubElement(annotation, 'size')
  ET.SubElement(size, 'width').text = str(width)
  ET.SubElement(size, 'height').text = str(height)
  ET.SubElement(size, 'depth').text = '3'
  ET.SubElement(annotation, 'segmented').text = '0'

  # 读取对应的txt文件并添加object信息
  with open(txt_path, 'r') as file:
    for line in file:
      parts = line.strip().split()
      obj = ET.SubElement(annotation, 'object')
      ET.SubElement(obj, 'name').text = classes[int(parts[0])]
      ET.SubElement(obj, 'pose').text = 'Unspecified'
      ET.SubElement(obj, 'truncated').text = '0'
      ET.SubElement(obj, 'difficult').text = '0'
      bndbox = ET.SubElement(obj, 'bndbox')
      # YOLO 格式转换为 VOC
      x_center = float(parts[1]) * width
      y_center = float(parts[2]) * height
      w = float(parts[3]) * width
      h = float(parts[4]) * height
      x_min = int(x_center - w / 2)
      y_min = int(y_center - h / 2)
      x_max = int(x_center + w / 2)
      y_max = int(y_center + h / 2)
      ET.SubElement(bndbox, 'xmin').text = str(x_min)
      ET.SubElement(bndbox, 'ymin').text = str(y_min)
      ET.SubElement(bndbox, 'xmax').text = str(x_max)
      ET.SubElement(bndbox, 'ymax').text = str(y_max)

  prettify(annotation)
  # 生成XML树并保存
  tree = ET.ElementTree(annotation)
  tree.write(xml_path)


def convert_yolo_to_voc(images_folder, labels_folder, output_folder, classes):
  # 创建输出文件夹
  if not os.path.exists(output_folder):
    os.makedirs(output_folder)

  # 遍历所有图像和标注
  for filename in os.listdir(images_folder):
    if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
      img_path = os.path.join(images_folder, filename)
      txt_filename = os.path.splitext(filename)[0] + '.txt'
      txt_path = os.path.join(labels_folder, txt_filename)
      xml_filename = os.path.splitext(filename)[0] + '.xml'
      xml_path = os.path.join(output_folder, xml_filename)
      if os.path.exists(txt_path):
        create_xml_file(img_path, txt_path, xml_path, classes)


# 用法示例
if __name__ == "__main__":
  images_folder = r'D:\file\postgrad\experiment\silver_spare_dataset\merge_labelframes_all1'
  labels_folder = r'D:\file\postgrad\experiment\silver_spare_dataset\merge_label_txt1'
  output_folder = r'D:\file\postgrad\experiment\silver_spare_dataset\merge_label_xml1'
  classes = ['asiaticblack bear', 'hare', 'hog badger', 'macaques', 'masked civet', 'muntjac', 'silver pheasant', 'squirrel', 'tragopan']  # 改成自己的类别

  convert_yolo_to_voc(images_folder, labels_folder, output_folder, classes)
